package com.itheima.ai.controller;

import com.itheima.ai.repository.ChatRepository;
import groovy.util.logging.Slf4j;
import lombok.RequiredArgsConstructor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.Objects;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;

@Slf4j
@CrossOrigin("*")
@RequiredArgsConstructor
@RestController
@RequestMapping("/ai/pdf")
public class PdfController {

    private static final Logger log = LoggerFactory.getLogger(PdfController.class);
    private final VectorStore vectorStore;

    private final ChatClient pdfChatClient;

    private final ChatRepository chatRepository;

    private final List<String> docIds;

    @PostMapping("/upload")
    public String upload(@RequestParam("file") MultipartFile file) {
        // 1. 校验文件是否为PDF格式
        if (!Objects.equals(file.getContentType(), "application/pdf")) {
            return "{\"ok\": 0}";
        }

        // 2.先清空旧文档向量
        vectorStore.delete(docIds);

        try {
            // 3. 使用PagePdfDocumentReader读取文件为Document格式
            PagePdfDocumentReader reader = new PagePdfDocumentReader(file.getResource(), PdfDocumentReaderConfig.builder()
                    .withPageTopMargin(0)
                    .withPageExtractedTextFormatter(ExtractedTextFormatter.builder()
                            .withNumberOfTopTextLinesToDelete(0)
                            .build())
                    .withPagesPerDocument(1)
                    .build());
            List<Document> documents = reader.read();
            // 4. 将文档写入向量数据库
            vectorStore.add(documents);
            // 5. 保存文档id，用于将来清空文档
            docIds.addAll(documents.stream().map(Document::getId).toList());
            return "{\"ok\": 1}";
        } catch (Exception e) {
            log.error("Failed to upload PDF file.", e);
            return "{\"ok\": 0}";
        }
    }

    @RequestMapping(value = "/chat", produces = "text/html;charset=UTF-8")
    public Flux<String> chat(String prompt, String chatId) {
        chatRepository.addChatId("pdf", chatId);
        return pdfChatClient
                .prompt(prompt)
                .advisors(a -> a.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId))
                .stream()
                .content();
    }

}
